Skip to content

Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules#4292

Merged
buptzyb merged 8 commits into
NVIDIA:mainfrom
buptzyb:main-strict-refactor
May 18, 2026
Merged

Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules#4292
buptzyb merged 8 commits into
NVIDIA:mainfrom
buptzyb:main-strict-refactor

Conversation

@buptzyb
Copy link
Copy Markdown
Contributor

@buptzyb buptzyb commented Apr 14, 2026

What does this PR do ?

Dev PR #4293

This PR decomposes the overloaded cuda_graph_scope field into three dedicated, semantically distinct concepts and cleans up naming throughout.

Problem

The old API overloaded --cuda-graph-scope with three unrelated concerns:

  • --cuda-graph-scope full_iteration → full-iteration training graphs (a capture strategy, not a module)
  • --cuda-graph-scope full_iteration_inference → block-owned inference graphs (inference ownership, not a module)
  • --cuda-graph-scope attn mlp → per-layer capture regions (the actual intended use)

CudaGraphScope mixed iteration-level control flow with per-layer module selection. The three concepts have nothing in common and cannot be meaningfully combined in one field.

Solution

Four concrete changes:

  1. full_iteration becomes its own cuda_graph_impl value.
    --cuda-graph-impl full_iteration replaces --cuda-graph-impl local --cuda-graph-scope full_iteration.

  2. --inference-cuda-graph-scope is a new dedicated field.
    InferenceCudaGraphScope (none / layer / block) replaces full_iteration_inference in cuda_graph_scope. The default for --cuda-graph-impl local is layer (preserving prior behaviour).

  3. CudaGraphScope is renamed to CudaGraphModule; cuda_graph_scope to cuda_graph_modules.
    With the two non-module values removed, the enum and field names now accurately reflect their purpose: selecting which per-layer modules to capture.

  4. Normalization is centralized in cuda_graph_config.py.
    normalize_cuda_graph_modules, normalize_inference_cuda_graph_scope, and validate_deprecated_cuda_graph_modules_migration_inputs are shared between TransformerConfig.__post_init__ and validate_args, removing duplication and making the migration logic testable in isolation.

Backward compatibility

All deprecated inputs are still accepted and silently migrated at startup:

Old input New equivalent
--enable-cuda-graph --cuda-graph-impl local
--external-cuda-graph --cuda-graph-impl transformer_engine
--cuda-graph-scope full_iteration --cuda-graph-impl full_iteration
--cuda-graph-scope full_iteration_inference --cuda-graph-impl local --inference-cuda-graph-scope block
--cuda-graph-scope attn mlp --cuda-graph-modules attn mlp
from megatron.core.transformer.enums import CudaGraphScope CudaGraphScope alias kept; use CudaGraphModule
TransformerConfig(cuda_graph_scope=...) cuda_graph_scope field kept, deprecated; use cuda_graph_modules

Conflicting combinations (e.g. passing both a deprecated value and the new flag with an incompatible value) are rejected with a clear assertion.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests
  • I have added proper typing to my code Typing guidelines
  • I have added relevant documentation
  • I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

  1. When your PR is ready, click Ready for Review.
  2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
    • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented Apr 14, 2026

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@buptzyb buptzyb force-pushed the main-strict-refactor branch 2 times, most recently from f36afad to 0f07b30 Compare April 23, 2026 14:55
@buptzyb buptzyb self-assigned this Apr 23, 2026
@buptzyb buptzyb changed the title Refactor CUDA graph configuration: separate inference granularity from training scope Refactor CUDA graph configuration: rename cuda_graph_scope to cuda_graph_modules and add backward compat Apr 23, 2026
@buptzyb buptzyb changed the title Refactor CUDA graph configuration: rename cuda_graph_scope to cuda_graph_modules and add backward compat Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules Apr 23, 2026
@buptzyb buptzyb marked this pull request as ready for review April 23, 2026 15:11
@buptzyb buptzyb requested review from a team as code owners April 23, 2026 15:11
@svcnvidia-nemo-ci svcnvidia-nemo-ci requested a review from a team April 23, 2026 15:11
@buptzyb buptzyb force-pushed the main-strict-refactor branch from 0f07b30 to 113c46a Compare April 23, 2026 15:18
@buptzyb buptzyb force-pushed the main-strict-refactor branch 5 times, most recently from 614115f to 9461d5a Compare April 26, 2026 11:02
@svcnvidia-nemo-ci svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels May 15, 2026
@buptzyb buptzyb enabled auto-merge May 15, 2026 11:35
@buptzyb
Copy link
Copy Markdown
Contributor Author

buptzyb commented May 15, 2026

/ok to test 73b71e7

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25928397787

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25928461458

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25938098213

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25938124861

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25942161774

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25947850162

@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks May 16, 2026
@buptzyb buptzyb added this pull request to the merge queue May 16, 2026
@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25954504754

@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks May 16, 2026
@buptzyb buptzyb added this pull request to the merge queue May 16, 2026
@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25957305634

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25963582311

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25967578667

buptzyb and others added 3 commits May 16, 2026 18:26
…_impl API

PR NVIDIA#3797 (A2A overlap for Megatron-FSDP) landed on main after this refactor
branch diverged, and it added a partial-CUDA-graph guard in
mcore_fsdp_adapter.py that iterates the legacy config.cuda_graph_scope
list. After the merge, that field is normalized to None and the iteration
raises TypeError: 'NoneType' object is not iterable, breaking every
a2a_overlap unit test and the deepseek_proxy_fsdp_ep2_fsdp2_ep_overlap
functional test.

Replace the legacy iteration with the equivalent new-API assertion on
config.cuda_graph_impl: per-layer CUDA graphs (local / transformer_engine)
are still rejected; full_iteration and none remain allowed. Drop the
now-unused CudaGraphScope import.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@buptzyb
Copy link
Copy Markdown
Contributor Author

buptzyb commented May 17, 2026

/ok to test e08b7e2

@svcnvidia-nemo-ci
Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26007684237

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved All necessary approvals have been made complexity: high Run functional tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.